This looks good, and if it fixes the problem, may be good enough. I think there is also a change to invest_all itself that should happen, so that instead of returning a modified copy of the original state, it mutates the state. I suppose that's part of another pull request. Have you verified that this does fix the problem? It will still be working with a copy of the state, and not the original object.